package com.aiwiown.snackmq.network.handler;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.util.MessageConverter;
import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.EncoderException;
import io.netty.handler.codec.MessageToMessageCodec;
import lombok.extern.slf4j.Slf4j;

import java.nio.ByteBuffer;
import java.util.List;

/**
 * 【已重构】自定义消息编解码器，现在支持两种协议格式：
 * 1. 标准消息 (Protobuf 序列化)
 * 2. 零拷贝批量响应 (原始 ByteBuffer 列表)
 *
 * 它通过一个 "帧类型" 字节来区分这两种格式，确保了向后兼容性和新功能的高效实现。
 */
@Slf4j
@ChannelHandler.Sharable // 标记为可共享，因为它不包含有状态的实例变量
public class MessageCodec extends MessageToMessageCodec<ByteBuf, Object> { // 泛型改为 <ByteBuf, Object>

    // 协议幻数，用于快速识别和验证数据包
    private static final int MAGIC_NUMBER = 0xCAFEBABE;
    // 协议版本号，为未来扩展保留
    private static final byte PROTOCOL_VERSION = 1;

    // --- 新增：帧类型常量，用于区分不同的响应格式 ---
    /**
     * 标准消息帧：整个消息体是一个 Protobuf 序列化的 Message 对象。
     */
    private static final byte FRAME_TYPE_STANDARD_MESSAGE = 1;
    /**
     * 零拷贝批量帧：响应体包含一个请求ID和一系列原始消息的 ByteBuffer。
     */
    private static final byte FRAME_TYPE_ZERO_COPY_BATCH = 2;


    /**
     * 编码方法，将 Message 或 ZeroCopyResponse 对象编码为 ByteBuf。
     *
     * @param ctx ChannelHandlerContext
     * @param msg 要编码的对象，可以是 Message 或 NettyResponseWriter.ZeroCopyResponse
     * @param out 编码后的 ByteBuf 列表
     */
    @Override
    protected void encode(ChannelHandlerContext ctx, Object msg, List<Object> out) throws Exception {
        ByteBuf buffer = ctx.alloc().buffer();
        // 1. 写入公共头部：幻数和版本号
        buffer.writeInt(MAGIC_NUMBER);
        buffer.writeByte(PROTOCOL_VERSION);

        // 2. 根据消息类型选择不同的编码路径
        if (msg instanceof Message) {
            encodeStandardMessage((Message) msg, buffer);
        } else if (msg instanceof NettyResponseWriter.ZeroCopyResponse) {
            encodeZeroCopyResponse((NettyResponseWriter.ZeroCopyResponse) msg, buffer);
        } else {
            // 【核心修复】: 当遇到未知类型时，释放 buffer 并抛出异常，而不是静默返回。
            // 这遵循了 MessageToMessageCodec 的规范，并使错误能被上层捕获。
            buffer.release();
            throw new EncoderException("MessageCodec received an unknown type to encode: " + msg.getClass().getName());
        }
        out.add(buffer);
    }

    /**
     * 编码标准 Message 对象。
     * 协议格式: [幻数(4B)] [版本(1B)] [总长度(4B)] [帧类型(1B)] [Protobuf 数据(可变)]
     */
    private void encodeStandardMessage(Message message, ByteBuf buffer) {
        byte[] bodyBytes = MessageConverter.toBytes(message);
        if (bodyBytes == null) {
            // 如果序列化失败，写入一个空的长度，避免发送损坏的数据
            buffer.writeInt(1);
            buffer.writeByte(FRAME_TYPE_STANDARD_MESSAGE);
            log.warn("Failed to serialize standard message, it was null. Sending empty message.");
            return;
        }
        // 写入总长度 (帧类型字节 + Protobuf数据长度)
        buffer.writeInt(1 + bodyBytes.length);
        // 写入帧类型
        buffer.writeByte(FRAME_TYPE_STANDARD_MESSAGE);
        // 写入 Protobuf 数据
        buffer.writeBytes(bodyBytes);
    }

    /**
     * 编码零拷贝响应。
     * 协议格式: [幻数(4B)] [版本(1B)] [总长度(4B)] [帧类型(1B)] [请求ID(8B)] [消息数量(4B)] [消息1长度(4B)] [消息1数据] ...
     */
    private void encodeZeroCopyResponse(NettyResponseWriter.ZeroCopyResponse zcResponse, ByteBuf buffer) {
        List<ByteBuffer> buffers = zcResponse.getBuffers();
        if (buffers == null || buffers.isEmpty()) {
            // 写入一个空的长度
            buffer.writeInt(1 + 8 + 4);
            buffer.writeByte(FRAME_TYPE_ZERO_COPY_BATCH);
            buffer.writeLong(zcResponse.getMessageId());
            buffer.writeInt(0);
            return;
        }

        // 1. 计算所有原始消息数据的总负载大小
        int payloadSize = 0;
        for (ByteBuffer b : buffers) {
            // 每个消息的长度(4B) + 消息自身数据
            payloadSize += (4 + b.remaining());
        }

        // 2. 写入总长度 (帧类型 + 请求ID + 消息数量 + 负载)
        buffer.writeInt(1 + 8 + 4 + payloadSize);
        // 3. 写入帧类型
        buffer.writeByte(FRAME_TYPE_ZERO_COPY_BATCH);
        // 4. 写入请求ID
        buffer.writeLong(zcResponse.getMessageId());
        // 5. 写入消息数量
        buffer.writeInt(buffers.size());
        // 6. 依次写入每个消息的长度和内容
        for (ByteBuffer b : buffers) {
            buffer.writeInt(b.remaining());
            buffer.writeBytes(b);
        }
    }


    /**
     * 解码方法，将 ByteBuf 解码为 Message 对象。
     * 服务端主要处理来自客户端的请求，这些请求都是标准 Message 格式。
     * 客户端的解码器需要处理服务端的两种响应格式。
     */
    @Override
    protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
        // 1. 读取并校验幻数和版本号
        int magicNumber = in.readInt();
        if (magicNumber != MAGIC_NUMBER) {
            log.error("Invalid magic number: {}, closing connection.", Integer.toHexString(magicNumber));
            ctx.close();
            return;
        }
        in.readByte(); // 读取并忽略版本号，为未来扩展保留

        // 2. 读取总长度。注意：LengthFieldBasedFrameDecoder 已经帮我们处理了半包问题。
        //    这里的长度是 [帧类型 + 剩余数据] 的长度。
        int length = in.readInt();
        if (in.readableBytes() < length) {
            // 理论上不应该发生，因为 LengthFieldBasedFrameDecoder 会保证这一点
            log.error("Readable bytes ({}) is less than declared length ({}).", in.readableBytes(), length);
            ctx.close();
            return;
        }

        // 3. 读取帧类型，以决定如何解码
        byte frameType = in.readByte();
        switch (frameType) {
            case FRAME_TYPE_STANDARD_MESSAGE:
                // 减去帧类型占用的1个字节
                int bodyLength = length - 1;
                if (bodyLength > 0) {
                    byte[] bodyBytes = new byte[bodyLength];
                    in.readBytes(bodyBytes);
                    Message domainMessage = MessageConverter.fromBytes(bodyBytes);
                    if (domainMessage != null) {
                        out.add(domainMessage);
                    }
                }
                break;

            // 服务端解码器通常不需要处理零拷贝响应，因为这是服务端发往客户端的。
            // 但为了健壮性，我们还是加上处理逻辑。
            case FRAME_TYPE_ZERO_COPY_BATCH:
                log.warn("Server received a ZERO_COPY_BATCH frame, which is unexpected. Ignoring.");
                // 跳过这个帧的剩余部分
                in.skipBytes(length - 1);
                break;

            default:
                log.error("Unknown frame type: {}. Closing connection.", frameType);
                ctx.close();
                break;
        }
    }

    /**
     * 异常处理
     */
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        log.error("Exception caught in MessageCodec", cause);
        ctx.close();
        // 将异常传递给 pipeline 中的下一个处理器，这是测试能够捕获异常的关键
        super.exceptionCaught(ctx, cause);
    }
}